#!/bin/bash

## set _CONDOR_IHEP_JOB_ID
CONDOR_IHEP_JOB_ID=`grep -E '^GlobalJobId ' $_CONDOR_JOB_AD | awk -F'#' '{print $2}'`
export _CONDOR_IHEP_JOB_ID=$CONDOR_IHEP_JOB_ID

# Export _CONDOR_IHEP_REMOTE_HOST, user can get where the job is running on
CONDOR_IHEP_REMOTE_HOST=`/bin/grep -E '^RemoteHost ' $_CONDOR_JOB_AD | /bin/awk -F '[ "]' '{print $4}'`
export _CONDOR_IHEP_REMOTE_HOST="$CONDOR_IHEP_REMOTE_HOST"

# Export _CONDOR_IHEP_SUBMITTED_TIME, user can get when the job was submitted
CONDOR_IHEP_SUBMISSION_TIME=`/bin/grep -E '^QDate ' $_CONDOR_JOB_AD | /bin/awk '{print $3}'`
CONDOR_IHEP_SUBMISSION_TIME=`/bin/date -d @${CONDOR_IHEP_SUBMISSION_TIME} +"%Y-%m-%d %T"`
export _CONDOR_IHEP_SUBMISSION_TIME="$CONDOR_IHEP_SUBMISSION_TIME"

# Check if requesting singularity, result 0 indicates requesting SL6 singularity
request_singularity=`/bin/grep -E '^IHEP_RequestContainer ' $_CONDOR_JOB_AD | /bin/awk -F '"' '{print $2}'`
if [[ -n "$request_singularity" ]] && [[ "$request_singularity" != "SL6" && "$request_singularity" != "SL7" ]];
then
        /bin/echo "ERROR: No container '${request_singularity}' under serving." >&2
        exit 19
fi

# Check if requesting gpu
request_gpu=`/bin/grep -E '^RequestGPUs ' $_CONDOR_JOB_AD | /bin/awk '{print $3}'`

## check OS version of this node
sl_version=`sed -r 's/[^0-9]*([0-9]+).*/\1/' /etc/redhat-release`

## detect acctgroupuser and acctgroup
owner=`grep -E '^Owner ' $_CONDOR_JOB_AD |awk -F'"' '{print $2}'`
ownergroups=($(/usr/bin/groups $owner 2>&1))
unset ownergroups[0] ownergroups[1]
ownergroups=("${ownergroups[@]}" "short" "test" "long")

acctgroupuser=`grep -E '^AcctGroupUser ' $_CONDOR_JOB_AD |awk -F'"' '{print $2}'`
acctgroup=`grep -E '^AcctGroup ' $_CONDOR_JOB_AD |awk -F'"' '{print $2}'`

if [[ ${ownergroups[@]/$acctgroup/} == ${ownergroups[@]} ]];
then
        echo "FATAL: Please set your own group." >&2
        exit 20
fi

if [[ "$owner" != "$acctgroupuser" ]];
then
        echo "FATAL: Please set your own username." >&2
        exit 21

fi

## extract mem information
#TotalMem=`/usr/bin/free -k |grep Mem | awk '{print $2}'`
#TotalCpu=`/usr/bin/lscpu |grep "^CPU(s):" | awk '{print $2}'`

## The handler of signals
trap stop_user_job 2 3 15
stop_user_job() {
    #echo "Job ${_CONDOR_IHEP_JOB_ID} is stopped by signal"
    kill -TERM $JobPid
}

## execute user's job with maximum (JobLimit) wall time
JobWallTime=`grep -E '^WallTime ' $_CONDOR_JOB_AD |awk -F'"' '{print $2}'`
walltime_limit=360000  # 100 hours
if [[ "$JobWallTime" == "long" ]]; then
    walltime_limit=2592000  # 30 days
fi

# --------------------------------------------------------------------------- #
# singularity area
# --------------------------------------------------------------------------- #
image_sl6="/cvmfs/container.ihep.ac.cn/singularity/image/SL69/sl69worknode20200729.sif"
image_sl7="/cvmfs/container.ihep.ac.cn/singularity/image/SL7/sl7worknode20200729.sif"
#singularity_path="/usr/bin/singularity"
container_path="/cvmfs/container.ihep.ac.cn/bin/hep_container"
singularity_cmd="$container_path exec"

export APPTAINERENV_SING_USER_DEFINED_APPEND_PATH=$PATH
#export SINGULARITYENV_APPEND_PATH=$PATH
#export SINGULARITYENV_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
if [[ -n "$LD_LIBRARY_PATH" ]]; then
	export APPTAINERENV_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
fi

if [[ -n "$SHELL" ]] ;
then
        export APPTAINER_SHELL=$SHELL
else
        current_user=`/bin/whoami`
        user_shell=`/bin/awk -F ":" '{print $1 ":"$NF}' /etc/passwd|grep -w $current_user|/bin/awk -F ":" '{print $2}'`
        export APPTAINER_SHELL=$user_shell
fi

# --------------------------------------------------------------------------- #
# run job
# --------------------------------------------------------------------------- #
user_exe=$1
shift

# Check if a singularity is needed and pick a proper image
if [[ -z "$request_singularity" ]]; then
        /usr/bin/timeout $walltime_limit $singularity_cmd CentOS7 ${user_exe} "$@" &
elif [[ "$request_singularity" == "SL6" ]]; then
        #/usr/bin/timeout $walltime_limit $singularity_cmd ${image_sl6} ${user_exe} "$@" &
        /usr/bin/timeout $walltime_limit $singularity_cmd SL6 ${user_exe} "$@" &
elif [[ "$request_singularity" == "SL7" ]]; then
        #/usr/bin/timeout $walltime_limit $singularity_cmd ${image_sl7} ${user_exe} "$@" &
        /usr/bin/timeout $walltime_limit $singularity_cmd SL7 ${user_exe} "$@" &
else
        /usr/bin/timeout $walltime_limit $singularity_cmd CentOS7 ${user_exe} "$@" &
fi

## Capture the pid of the job to kill later, if needed
JobPid=$!
## Wait for the job to exit of own accord
wait $JobPid
ExitCode=$?
if [ $ExitCode -eq 124 ]; then
    echo "Job is KILLED: exceeds the 100 hours time limit" >&2
fi

## If we get here, the job exitted. Put code to run after job here...
post_script="$script_path/post.sh"
if [[ -f "$post_script" ]]; then
        /bin/bash $post_script $acctgroupuser $acctgroup
        exit_status=$?
        if [[ $exit_status -ne 0 ]]; then
                echo "ERROR: Your job can not be finalized" >&2
                echo "Please contact helpdesk@ihep.ac.cn." >&2
                exit 29
        fi
fi

## Exit with the jobs exit code
exit $ExitCode
